{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C2mxc11jnyOP"
   },
   "source": [
    "# Hierarchial Generation of Moelcules\n",
    "\n",
    "In this tutorial, we will load a pre-trained model which hierarchically generates molecules using graph motifs. Prior work on generating molecules tends to use small units such as atoms and simpler cycles, but this limits the graph generation to smaller molecules. In Hierarchical Generation of Molecular Graphs using Structural Motifs (https://arxiv.org/abs/2002.03230), Jin et. al. proposed a method to use larger and more flexible motifs as units. \n",
    "\n",
    "Motifs are produced seperately before the generation process based on frequenclty occuring substructures. To create motifs, a molecule is decomposed into fragments by breaking bridge bonds that wouldn't violate chemical validity. A fragment is selected as a motif if it occurs in the training set more than a set amount of times. If the fragment is not selected, then break them into rings and bonds and select as a motif in the original molecule. This allows for the creating of a vocabulary of motifs.\n",
    "\n",
    "The molecular generation model is an extension of the VAE by using a hierarchial encoder and decoder. During the generation step, the decoder adds a motif and attachment cofiguration and then decides how it should be added to the current graph. Thus, the encoder needs to learn to represent molecules at different scales. So, a molecule is represented as a hiearchial graph with 3 levels: motif, attachment, and atom. There are also edges connecting atoms and motifs between different layers and by using a hiearchial message passing network (MPN), the molecule is encoded.\n",
    "\n",
    "![alt text]()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RDKit Version\n",
    "\n",
    "An important caveat is that this code does not work with the latest version of RDKit. Currently, it is known that it works with RDKit version >= 2019.03 and <= 2020.09.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jcYi99MDnyOW"
   },
   "source": [
    "## Colab\n",
    "This tutorial and the rest in this sequence are designed to be done in Google colab. If you'd like to open this notebook in colab, you can use the following link.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deepchem/deepchem/blob/master/examples/tutorials/Generating_molecules_with_MolGAN.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1VTHRshpnyOX"
   },
   "source": [
    "## Setup\n",
    "\n",
    "The original repository had to be modified to allow running on a cpu, so we are using a fork"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "E8kXK8sunyOa",
    "outputId": "73c3da20-30a5-4abc-8dde-6775e7ffd06f"
   },
   "outputs": [],
   "source": [
    "%cd drive/MyDrive\n",
    "!git clone https://github.com/saithat/hgraph2graph.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "SUPYyPENnyOd",
    "outputId": "993463d8-c1f7-49bd-a13b-77fbf9302c1b"
   },
   "outputs": [],
   "source": [
    "%cd hgraph2graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 341
    },
    "id": "5J9_6xShpgbf",
    "outputId": "ed1fcb21-56ec-42e6-d308-252ed070915a"
   },
   "outputs": [],
   "source": [
    "!pip install --pre deepchem\n",
    "import deepchem\n",
    "deepchem.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "RhbkFQdosQ01"
   },
   "outputs": [],
   "source": [
    "#!pip install chemprop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "25XxHiPknyOe"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import os\n",
    "from collections import OrderedDict\n",
    "\n",
    "import rdkit\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem.Draw import IPythonConsole\n",
    "from rdkit.Chem import Draw\n",
    "\n",
    "\n",
    "import math, random, sys\n",
    "import numpy as np\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "Fiu9CQUBnyOg"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "6YK1ulbcnyOi"
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'hgraph'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-95393c9e2a75>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#import finetune_generator\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;31m#from finetune_generator import Chemprop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mhgraph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mhgraph\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'hgraph'"
     ]
    }
   ],
   "source": [
    "#import finetune_generator\n",
    "#from finetune_generator import Chemprop\n",
    "import hgraph\n",
    "from hgraph import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R_Yb8UFWnyOl"
   },
   "source": [
    "Original repository uses extensive command line arguments, so to modify as least as possible, we define our own Args class to fit into existing structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RPdHLwy7nyOp"
   },
   "outputs": [],
   "source": [
    "class Args():\n",
    "    def __init__(self):\n",
    "        self.vocab = 'data/chembl/vocab.txt'\n",
    "        self.atom_vocab = common_atom_vocab\n",
    "        self.model = 'ckpt/chembl-pretrained/model.ckpt'\n",
    "        self.seed = 7\n",
    "        self.nsample = 100\n",
    "        self.rnn_type = 'LSTM'\n",
    "        self.hidden_size = 250\n",
    "        self.embed_size = 250\n",
    "        self.batch_size = 50\n",
    "        self.latent_size = 32\n",
    "        self.depthT = 15\n",
    "        self.depthG = 15\n",
    "        self.diterT = 1\n",
    "        self.diterG = 3\n",
    "        self.dropout = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "43x1cZThnyOr"
   },
   "outputs": [],
   "source": [
    "args = Args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZOOXseAUnyOt"
   },
   "source": [
    "Open the pre-defined motif vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EZXiZOcinyOu"
   },
   "outputs": [],
   "source": [
    "vocab = [x.strip(\"\\r\\n \").split() for x in open(\"data/chembl/vocab.txt\")]\n",
    "args.vocab = PairVocab(vocab, cuda=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S5Q_ip5fnyOw"
   },
   "source": [
    "Code in existing repo was designed to work with Nvidia GPUs, so we modified to allow working with cpu as well in case a Nvidia GPU is not available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "w_ekBTxNnyOx",
    "outputId": "50d450ac-f6f2-417b-de12-3a3cb36a9115"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NGyisMGgnyOy"
   },
   "source": [
    "Load pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qoczA7mUnyOz",
    "outputId": "c85a642a-e203-4652-9a86-a98bf7172ded"
   },
   "outputs": [],
   "source": [
    "model = HierVAE(args).to(device)\n",
    "if not torch.cuda.is_available() == 'cpu':\n",
    "    model.load_state_dict(torch.load(args.model, map_location=torch.device('cpu'))[0])\n",
    "else:\n",
    "    model.load_state_dict(torch.load(args.model)[0])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YupnqdianyO2"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(args.seed)\n",
    "random.seed(args.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C8WjO5SFnyO3"
   },
   "outputs": [],
   "source": [
    "nmols_smiles = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OYNG4rv_nyO4"
   },
   "source": [
    "Sample molecules from the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 391
    },
    "id": "DeANuVVnnyO5",
    "outputId": "78cec8df-852e-449e-c931-668e149a35f3"
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    for _ in tqdm(range(args.nsample // args.batch_size)):\n",
    "        smiles_list = model.sample(args.batch_size, greedy=True)\n",
    "        for _,smiles in enumerate(smiles_list):\n",
    "            print(smiles)\n",
    "            nmols_smiles.append(smiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gyclFUk9nyO7"
   },
   "source": [
    "Keep only the unique molecules and convert to Rdkit Mols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F3uf-vr0nyO8"
   },
   "outputs": [],
   "source": [
    "nmols_smiles_unique = list(OrderedDict.fromkeys(nmols_smiles))\n",
    "nmols_viz = [Chem.MolFromSmiles(x) for x in nmols_smiles_unique]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OJnNKXXLnyO9"
   },
   "outputs": [],
   "source": [
    "img = Draw.MolsToGridImage(nmols_viz[0:100], molsPerRow=5, subImgSize=(250, 250), maxMols=100, legends=None, returnPNG=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JlgXbgL-nyO-"
   },
   "outputs": [],
   "source": [
    "img"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Hierarchial_Moelcular_Generation_with_HierVAE.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "deepchem",
   "language": "python",
   "name": "deepchem"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
